import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk, ImageFilter
import numpy as np
from util import preprocess_image, load, get

char_styles = get('char_styles')
new_size = get('new_size')

class ImageClassifierApp:
    def __init__(self, model_path):
        self.model = load("F1最高模型", f'{get("model_root")}/best_model')

        self.root = tk.Tk()
        self.root.title('Image Classifier')
        self.root.geometry("800x600")

        self.button = tk.Button(self.root, text='选择图像', command=self.select_image)
        self.button.pack()

        self.image_label = tk.Label(self.root)
        self.image_label.pack()

        self.prediction_label = tk.Label(self.root)
        self.prediction_label.pack()

        # 添加一个新的标签用于显示书法类别
        self.style_label = tk.Label(self.root)
        self.style_label.pack()

        self.root.mainloop()

    def select_image(self):
        image_path = filedialog.askopenfilename()

        img_test = np.array(preprocess_image(image_path, new_size)).reshape(1,-1)

        predicted_class = self.model.predict(img_test)

        pil_image = Image.open(image_path)
        pil_image = pil_image.resize((800, 600), Image.Resampling.LANCZOS)

        image_tk = ImageTk.PhotoImage(pil_image)
        self.image_label.config(image=image_tk)
        self.image_label.image = image_tk

        self.prediction_label.config(text=f'预测类别: {char_styles[predicted_class[0]]}')

        # 更新书法类别标签
        self.style_label.config(text=f'书法类别: {char_styles[predicted_class[0]]}')

app = ImageClassifierApp(f'{get("model_root")}/best_model')
